package z1crypto

import (
	"bytes"
	"crypto/aes"
	"crypto/cipher"
	"crypto/md5"
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"errors"
	"strings"

	"encoding/base64"
	"encoding/hex"
	"encoding/pem"

	"gitee.com/z1gotool/z1err"
)

func AesEncrypt(plaintext, key string, ivs ...string) (ret string, err error) {
	z1err.Check(err)
	// https://news.sangniao.com/p/3723057230
	// https://blog.csdn.net/weixin_41806099/article/details/132137943

	iv := fixIv(ivs...)

	var plainTextBlock []byte
	length := len(plaintext)

	if length%16 != 0 {
		extendBlock := 16 - (length % 16)
		plainTextBlock = make([]byte, length+extendBlock)
		copy(plainTextBlock[length:], bytes.Repeat([]byte{uint8(extendBlock)}, extendBlock))
	} else {
		plainTextBlock = make([]byte, length)
	}

	copy(plainTextBlock, plaintext)
	block, err := aes.NewCipher([]byte(key))
	z1err.Check(err)

	ciphertext := make([]byte, len(plainTextBlock))
	mode := cipher.NewCBCEncrypter(block, []byte(iv))
	mode.CryptBlocks(ciphertext, plainTextBlock)

	str := base64.StdEncoding.EncodeToString(ciphertext)

	ret = str

	return
}

func AesDecrypt(encrypted, key string, ivs ...string) (ret string, err error) {
	z1err.Handle(&err)
	// https://news.sangniao.com/p/3723057230

	iv := fixIv(ivs...)

	if len(encrypted) == 0 {
		return
	}

	ciphertext, err := base64.StdEncoding.DecodeString(encrypted)
	z1err.Check(err)

	block, err := aes.NewCipher([]byte(key))
	z1err.Check(err)

	if len(ciphertext)%aes.BlockSize != 0 {
		z1err.Check(errors.New("block size cant be zero"))
	}

	mode := cipher.NewCBCDecrypter(block, []byte(iv))
	mode.CryptBlocks(ciphertext, ciphertext)
	ciphertext = PKCS5UnPadding(ciphertext)
	ret = string(ciphertext)

	return
}

// PKCS5UnPadding用必要的数据来填充某个数据团，以用于AES块状密码。
func PKCS5UnPadding(src []byte) []byte {
	length := len(src)
	unpadding := int(src[length-1])

	if length >= unpadding {
		return src[:(length - unpadding)]
	} else {
		return src
	}
}

func fixIv(ivs ...string) (iv string) {
	if len(ivs) > 0 {
		iv = ivs[0]
	}

	iv = iv + `1234567890123456`

	iv = iv[0:16]

	return
}

func RsaKeyClean(key string) (base64 string) {

	infoStr := key

	infoStr = strings.ReplaceAll(infoStr, "\t", "")

	base64 = infoStr

	return
}

func RsaEncrypt(origDataStr, publicKeyStr string) (ciphertextBase64 string, err error) {
	z1err.Handle(&err)
	// https://www.yii666.com/blog/129530.html
	// https://blog.csdn.net/weixin_41806099/article/details/132137943

	publicKeyStr = RsaKeyClean(publicKeyStr)

	origData := []byte(origDataStr)
	publicKey := []byte(publicKeyStr)

	block, _ := pem.Decode(publicKey)
	if block == nil {
		z1err.Check(errors.New("public key error"))
	}

	pubInterface, err := x509.ParsePKIXPublicKey(block.Bytes)
	z1err.Check(err)
	pub := pubInterface.(*rsa.PublicKey)

	b, err := rsa.EncryptPKCS1v15(rand.Reader, pub, origData)
	z1err.Check(err)

	ciphertextBase64 = base64.StdEncoding.EncodeToString(b)

	return
}

func RsaDecrypt(ciphertextBase64, privateKeyStr string) (origData string, err error) {
	z1err.Handle(&err)
	// https://www.yii666.com/blog/129530.html

	privateKeyStr = RsaKeyClean(privateKeyStr)

	privateKey := []byte(privateKeyStr)
	ciphertext, err := base64.StdEncoding.DecodeString(ciphertextBase64)
	z1err.Check(err)

	block, _ := pem.Decode(privateKey)
	if block == nil {
		z1err.Check(errors.New("private key error!"))
	}
	priv, err := x509.ParsePKCS1PrivateKey(block.Bytes)
	z1err.Check(err)

	b, err := rsa.DecryptPKCS1v15(rand.Reader, priv, ciphertext)
	z1err.Check(err)

	origData = string(b)

	return
}

func Base64Encode(msg string, inUrl ...bool) (base64Str string) {
	url := false
	if len(inUrl) > 0 {
		url = inUrl[0]
	}

	if url {
		base64Str = base64.URLEncoding.EncodeToString([]byte(msg))
	} else {
		base64Str = base64.StdEncoding.EncodeToString([]byte(msg))
	}

	return
}

func Base64Decode(base64Str string, inUrl ...bool) (msg string, err error) {
	url := false
	if len(inUrl) > 0 {
		url = inUrl[0]
	}

	if url {
		dec_str, err1 := base64.URLEncoding.DecodeString(base64Str)
		msg = string(dec_str)
		err = err1
	} else {
		dec_str, err1 := base64.StdEncoding.DecodeString(base64Str)
		msg = string(dec_str)
		err = err1
	}

	return
}

func Md5(str string) (encode string) {
	md5 := md5.New()
	md5.Write([]byte(str))
	return hex.EncodeToString(md5.Sum(nil))
}
